import time
from fastapi import Depends, HTTPException, WebSocket
from asyncio import Queue,wait_for
from typing import Annotated

from loguru import logger

class ConnectionManager:
    def __init__(self):
        self.active_connections: dict[int, WebSocket] = {}
        self.active_web_connections: dict[int, WebSocket] = {}
        self.response_data: dict[str, Queue] = {}
    def add_queue(self, req_id:str):
        self.response_data[req_id] = Queue()
    async def add_data(self, req_id:str,data: dict):
        await self.response_data[req_id].put(data)
    def add(self, websocket: WebSocket, device_id: int):
        self.active_connections[device_id] = websocket
    def add_web(self, websocket: WebSocket, device_id: int):
        self.active_web_connections[device_id] = websocket

    async def get_resp_data(self, req_id:int,timeout:float=2)->dict:
        try:
            data = await wait_for(self.response_data[req_id].get(), timeout)
        except:
            raise HTTPException(status_code=408, detail="Request timed out")
        self.response_data[req_id].task_done()
        return data
            

    async def broadcast(self, message: dict):
        """向所有客户端广播消息"""
        for device_id, connection in list(self.active_connections.items()):
            try:
                await connection.send_json(message)
            except Exception as e:
                # 如果发送失败（如客户端已断开），移除连接
                logger.warning(f"Failed to send message: {e}")
                connection.close()
                self.active_connections.pop(device_id, None)

    async def broadcast_to(self, device_id: int, message: dict):
        """向指定客户端广播消息"""
        connection = self.active_connections.get(device_id)
        logger.info(f"Broadcasting to device {device_id}")
        logger.info(f"Message: {message}")
        logger.info(f"Connection: {connection}")
        try:
            await connection.send_json(message)
        except Exception as e:
            # 如果发送失败（如客户端已断开），移除连接
            logger.warning(f"Failed to send message: {e}")
    async def broadcast_to_web(self, message: dict):
        """向指定客户端广播消息"""
        for device_id, connection in list(self.active_web_connections.items()):
            try:
                await connection.send_json(message)
            except Exception as e:
                # 如果发送失败（如客户端已断开），移除连接
                logger.warning(f"Failed to send message: {e}")
                connection.close()
                self.active_web_connections.pop(device_id, None)


# 全局实例
manager = ConnectionManager()

def get_manager() -> ConnectionManager:
    return manager

ManagerDep = Annotated[ConnectionManager, Depends(get_manager)]